"""A parser for Newick (New Hampshire) format."""

from pyparsing import *
from phylopy.tree import Tree, Node

# Enable packrat parsing
ParserElement._parse = ParserElement._parseCache

# Labels
label = Word(alphanums + '_')

# Distances
distance = Suppress(":") + Word(nums + "+-.eE")
distance.setParseAction(lambda t: float(t[0]))

node = Forward()
children = Suppress("(") + delimitedList(node) + Suppress(")")

nodeC = children("children") + Optional(label)("label") + Optional(distance)("distance")
nodeL = Optional(children)("children") + label("label") + Optional(distance)("distance")
nodeD = Optional(children)("children") + Optional(label)("label") + distance("distance")

node << Group(nodeC | nodeL | nodeD)

tree = node + Suppress(";") + stringEnd


def processNodes(g, parent, res):
    """Recurses down the parse tree, creating Node objects and hooking them up to other graph
    elements with edges, as appropriate.  Note that recursion is probably safe here.  To reach
    Python's default recursion limit (99), one would have to have a tree with at least 
    exp(99 * ln(2)) leaves.  That's approaching the number of atoms in the universe, so if my
    calculations are correct, I think we are safe."""
    for elt in res:
        n = Node(elt.label)
        g.add_node(n)
        g.add_edge(parent, n, weight=elt.distance)
        if elt.children:
            processNodes(g, n, elt.children)


def makeTree(s):
    res = tree.parseString(s)
    g = Tree()
    res = res[0]
    parent = Node(res.label)
    g.add_node(parent)
    processNodes(g, parent, res.children)
    return g 


def runTests():
    """A cursory test harness."""
    tests = """\
        dog;
        (dog, cat, mouse)mammal;
        (dog, cat, (mouse, rat)rodent, gorilla)mammal;
        dog:10.1;
        ((dog:3, cat:7):4, mouse:8);""".splitlines()
    for t in tests:
        print t.strip() 
        t = makeTree(t)
        print t, [e.label for e in t.leaves()]

def testTree(treeFile):
    f = open(treeFile, 'r')
    t = makeTree(f.read())
    f.close()
    print len(t.leaves())
